/* $Id: communicate.cc jjs $ */
/*
   Copyright (C) 2007 by Jurjen Stellingwerff <jurjen@stwerff.xs4all.nl>
   Part of the Moros Project http://moros.sourceforge.net/

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License version 2 or any
   later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY.

   See the COPYING file for more details.
*/
#define DEBUG_ROUTINE "communicate"

#include "communicate.h"
#include "signal.h"
#include "util.h"
#include "log.h"
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <errno.h>
#include <time.h>
#include <netdb.h>
#include <iostream>
#include <netinet/in.h>
#include <arpa/inet.h>

#define MAXDATA 10240
#define DH_BITS 1024
#define dfile(filename) (data_dir+"/"+filename).c_str()


int tls_init=false;

const int cert_type_priority[2] = { GNUTLS_CRT_X509, 0 };
const int kx_prio[] = { GNUTLS_KX_DHE_DSS, 0 };
const int mac_prio[] = { GNUTLS_MAC_SHA1, 0 };

/// Read and write files as connection

class file_conn: public connection {
private:
	int fd;
public:
	file_conn(int fptr) {
		fd=fptr;
	}

	int read(char *buffer, size length) {
	    dbg_calls("read length=%u", length);
		int res=::read(fd, buffer, length);
		int err=errno;
		if (res<0 && errno==EAGAIN) return 0;
		if (res<0) dbg_state("Reading error: %s (%i)", strerror(err), err);
		dbg_data("Read %d bytes", res);
		return res;
	}

	int write(const char *buffer, size length) {
		int res=::write(fd, buffer, length);
		int err=errno;
		if (res<0 && errno==EAGAIN) return 0;
		if (res<0) dbg_state("Writng error: %s (%i)", strerror(err), err);
		dbg_data("Wrote %d bytes", res);;
		return res;
	}

	void close() {
		::close(fd);
	}

	size length() {
		size l=0;
		struct stat buf;
		if (fstat(fd, &buf) >= 0) {
			l = buf.st_size;
  		}
		return l;
	}
};

/// Read and write data over a normal internet socket

class socket_conn: public connection {
private:
	int sock;
public:
	socket_conn(int socket) {
		sock=socket;
	}
	int read(char *buffer, size length) {
		int res=recv(sock, buffer, length, MSG_DONTWAIT);
		if (res<0) {
			int err=errno;
			if (errno==EAGAIN) return 0;
			dbg_comm("Reading ERROR %s", strerror(err));
		} else if (res>0) dbg_data("Reading %i bytes", res);
		return res;
	}
	int write(const char *buffer, size length) {
		int res=send(sock, buffer, length, MSG_DONTWAIT);
		if (res<0) {
			int err=errno;
			if (errno==EAGAIN) return 0;
			dbg_comm("Writing ERROR %s", strerror(err));
		} else if (res>0) dbg_data("Writing %i bytes", res);
		return res;
	}
	void close() {
		shutdown(sock, SHUT_RDWR);
		::close(sock);
	}
	size length() {
		return -1;
	}
};

/// Read and write data over a secured socket

class tls_conn: public connection {
private:
	gnutls_session_t session;
	int sock;
public:
	tls_conn(int socket, gnutls_session_t tls_session) {
		sock=socket;
		session=tls_session;
	}
	int read(char *buffer, size length) {
		int res=gnutls_record_recv(session, buffer, length);
		if (res>0) {
		    buffer[res]=0;
		    return res;
        }
		if (res==GNUTLS_E_AGAIN || res==GNUTLS_E_INTERRUPTED) return 0;
		dbg_comm("Reading ERROR %s", gnutls_strerror(res));
		return res;
	}
	int write(const char *buffer, size length) {
		int res=gnutls_record_send(session, buffer, length);
		if (res==GNUTLS_E_AGAIN || res==GNUTLS_E_INTERRUPTED) return 0;
		if (res<0) dbg_comm("Writing ERROR %s", gnutls_strerror(res));
		return res;
	}
	void close() {
   	    dbg_comm("Closing connection");
		gnutls_bye (session, GNUTLS_SHUT_WR);
		shutdown(sock, SHUT_RDWR);
		::close(sock);
		gnutls_deinit (session);
	}
	size length() {
		return -1;
	}
};

/// Set the given file descriptor to block or unblock
#ifndef _WIN32

#include <fcntl.h>
void set_blocking (int fileno, bool blocking) {
    int flag;
    flag = fcntl (fileno, F_GETFL, 0);
    if (blocking) {
        flag &= (~O_NDELAY);
    } else {
        flag |= (O_NDELAY);
    }
    fcntl (fileno, F_SETFL, flag);
}

inline bool is_nonblocking_error() {
    switch (errno) {
      case EAGAIN:
      case EALREADY:
      case EINPROGRESS:
        return true;
    }
    return false;
}

#else // _WIN32

void set_blocking (int fileno, bool blocking) {
    u_long b = blocking;
    ioctlsocket(fileno, FIONBIO, &b);
}

bool is_nonblocking_error() {
    switch (WSAGetLastError()) {
      case WSAEWOULDBLOCK:
      case WSAEALREADY:
      case WSAEINPROGRESS:
        return true;
    }
    return false;
}
#endif // _WIN32

communicate::communicate() {
	signal(SIGPIPE, SIG_IGN);
	serv_socket_ = -1;
	channelnr = 1;
	type = 1;
	init = false;
	session="";
	sessnr=1;
	data_dir="/var/lib/moros";
}
communicate::~communicate() {
    #ifndef no_gnutls
	if (type==2) {
		gnutls_certificate_free_credentials(cred);
	}
	#endif
}

void communicate::set_dir(std::string dir) {
    data_dir=dir;
}

channel* communicate::reading(const std::string file, sigc::signal<void, channel*> handler) {
    dbg_calls("Reading file '%s'", file.c_str());
	int fd = open(file.c_str(), O_RDONLY + O_NONBLOCK);
	if (fd == -1) {
		int err=errno;
		dbg_error("Error %s reading file '%s'", strerror(err), file.c_str());
		return 0;
	}
	channel res=channel(2, 0);
	res.id=0; // special for files
	res.handler=handler;
	res.stream=new file_conn(fd);
	conns_.push_back(res);
	return &(*conns_.rbegin());
}

channel* communicate::writing(const std::string file) {
    dbg_calls("Writing file '%s'", file.c_str());
	int fd = open(file.c_str(), O_WRONLY + O_CREAT + O_TRUNC, S_IRUSR + S_IWUSR);
	if (fd == -1) {
		int err=errno;
		dbg_error("Error %s creating file '%s'", strerror(err), file.c_str());
		return 0;
	}
	channel res=channel(2, 0);
	res.id=0; // special for files
	res.stream=new file_conn(fd);
	conns_.push_back(res);
	return &(*conns_.rbegin());
}

bool communicate::reconnect(channel* ch) {
	struct hostent *h;
	struct sockaddr_in serv_addr;
	if ((h = gethostbyname(ch->ip.c_str()))==NULL) {
		dbg_comm("Cannot determin host %s", ch->ip.c_str());
		return false;
	}
	memset(&serv_addr, 0, sizeof(serv_addr));
	serv_addr.sin_family = AF_INET;
	serv_addr.sin_addr.s_addr = ((struct in_addr *)(h->h_addr))->s_addr;
	serv_addr.sin_port = htons(ch->port);
	int sock = socket(AF_INET, SOCK_STREAM, 0);
	if (sock<0) {
		dbg_comm("Error creating socket");
		return false;
  	}
	ch->stream=new socket_conn(sock);
	int optval=1;
	setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
  	if (connect(sock, (struct sockaddr *) &serv_addr, sizeof(serv_addr))<0) {
    	dbg_comm("Cannot connect to %s port %i", ch->ip.c_str(), ch->port);
    	return false;
  	}
  	ch->id=channelnr++;
  	ch->finished=false;
	dbg_comm("Opened channel %i", ch->id);
	if (session!="") {
		ch->send_session(time(NULL));
		ch->lock=true;
	}
	return true;
}

channel* communicate::client(const std::string host, int port, sigc::signal<void, channel*> handler) {
	serv_socket_ = -1;
	channel res=channel(1, time(NULL));
	res.handler=handler;
	res.comm=this;
	res.ip=host;
	res.port=port;
	if (!reconnect(&res)) return 0;
	conns_.push_back(res);
	return &(*conns_.rbegin());
}

#ifndef no_gnutls
channel* communicate::tls_client(const std::string host, int port, sigc::signal<void, channel*> handler) {
	struct hostent *h;
	struct sockaddr_in serv_addr;
	if (!tls_init) {
		gnutls_global_init();
		tls_init=true;
	}
	if (!init) {
		gnutls_certificate_allocate_credentials (&cred);
		int ret=gnutls_certificate_set_x509_trust_file(cred, dfile("ca.pem"), GNUTLS_X509_FMT_PEM);
		if (ret<0) {
            dbg_error("Cannot read CA file %s", dfile("ca.pem"));
            return 0;
		}
		init=true;
	}
	int ret;
	gnutls_session_t session;
	ret=gnutls_init (&session, GNUTLS_CLIENT);
	if (ret<0) {
	    dbg_error("Cannot initiate tls connection");
        return 0;
	}
	gnutls_set_default_priority (session);
  	ret=gnutls_certificate_type_set_priority (session, cert_type_priority);
	if (ret<0) {
	    dbg_error("Cannot set priority on tls connection");
        return 0;
	}
	ret=gnutls_credentials_set (session, GNUTLS_CRD_CERTIFICATE, cred);
	if (ret<0) {
	    dbg_error("Cannot set credentials on tls connection");
        return 0;
	}
	serv_socket_ = -1;
	channel res=channel(1, time(NULL));
	res.handler=handler;
	res.id=channelnr++;
	if ((h = gethostbyname(host.c_str()))==NULL) {
		dbg_comm("cannot find host %s", host.c_str());
		return 0;
	}
	memset(&serv_addr, 0, sizeof(serv_addr));
	serv_addr.sin_family = AF_INET;
	serv_addr.sin_addr.s_addr = ((struct in_addr *)(h->h_addr))->s_addr;
	serv_addr.sin_port = htons(port);
	int sock = socket(AF_INET, SOCK_STREAM, 0);
	if (sock<0) {
	    dbg_comm("Error creating socket");
        return 0;
  	}
	int optval=1;
	setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
  	if (connect(sock, (struct sockaddr *) &serv_addr, sizeof(serv_addr))<0) {
    	dbg_comm("Cannot connect to %s port %i", host.c_str(), port);
    	return 0;
  	}
	gnutls_transport_set_ptr (session, (gnutls_transport_ptr_t) sock);
	ret = gnutls_handshake (session);
	if (ret < 0) {
    	dbg_error("Handshake has failed (%s)", gnutls_strerror (ret));
		close (sock);
		gnutls_deinit (session);
    	return 0;
	}
	res.stream=new tls_conn(sock, session);
	dbg_calls("Opened tls session %i", res.id);
	set_blocking(sock, false);
	conns_.push_back(res);
	return &(*conns_.rbegin());
}
#endif

void communicate::server(int port, sigc::signal<void, channel*> handler) {
	struct sockaddr_in serv_addr;
	serv_socket_ = socket(AF_INET, SOCK_STREAM, 0);
	if (serv_socket_<0)	dbg_fatal("creating socket");
	memset(&serv_addr, 0, sizeof(serv_addr));
	serv_addr.sin_family = AF_INET;
	serv_addr.sin_addr.s_addr = INADDR_ANY;
	serv_addr.sin_port = htons(port);
	int optval=1;
	setsockopt(serv_socket_, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
	if (bind(serv_socket_, (struct sockaddr*) &serv_addr, sizeof(serv_addr))<0)
	  	dbg_fatal("Cannot listen on port %i allready a server running?", port);
	serv_handler = handler;
	listen(serv_socket_, 200);
	set_blocking(serv_socket_, false);
}

#ifndef no_gnutls
void communicate::tls_server(int port, sigc::signal<void, channel*> handler) {
	if (!tls_init) {
		gnutls_global_init();
		tls_init=true;
	}
	if (!init) {
		int res;
		res=gnutls_certificate_allocate_credentials (&cred);
		if (res!=0) dbg_fatal("gnutls_certificate_allocate_credentials");
		res=gnutls_certificate_set_x509_trust_file(cred, dfile("ca.pem"), GNUTLS_X509_FMT_PEM);
		if (res<1) dbg_fatal("gnutls_certificate_set_x509_trust_file '%s'", dfile("ca.pem"));
		res=gnutls_certificate_set_x509_key_file(cred, dfile("cert.pem"), dfile("key.pem"), GNUTLS_X509_FMT_PEM);
		if (res<0) dbg_fatal("gnutls_certificate_set_x509_key_file public '%s' private '%s'",
                        dfile("cert.pem"), dfile("key.pem"));
        res=gnutls_dh_params_init (&dh_params);
        if (res<0) dbg_fatal("gnutls_dh_params_init");
        res=gnutls_rsa_params_init (&rsa_params);
        if (res<0) dbg_fatal("gnutls_rsa_params_init");

        tls_reload();
		gnutls_certificate_set_dh_params (cred, dh_params);
		gnutls_certificate_set_rsa_export_params(cred, rsa_params);
		init=true;
	}
	server(port, handler);
	type=2;
}

char *fd_load(int fd, size_t *psize, size_t maxsize) {
	struct stat sbuf;
	size_t size;
	char *p;
	if (fstat(fd, &sbuf) != 0) return NULL;
	size = sbuf.st_size;
	if (maxsize) size = (size<maxsize?size:maxsize);
	p = (char *)malloc(size+1);
	if (!p) return NULL;
	if (read(fd, p, size) != (long)size) {
		free(p);
		return NULL;
	}
	p[size] = 0;
	if (psize) *psize = size;
	return p;
}

char *file_load(const char *fname, size_t *size, size_t maxsize) {
	int fd;
	char *p;
	if (!fname || !*fname) return NULL;
	fd = open(fname,O_RDONLY);
	if (fd == -1) return NULL;
	p = fd_load(fd, size, maxsize);
	close(fd);
	return p;
}

void communicate::tls_reload() {
    dbg_calls("tls_reload");
	gnutls_datum_t parms;
	size_t size;
	int res;
	parms.data = (uint8_t *)file_load(dfile("dhfile"), &size, MAXDATA);
	if (!parms.data) {
	    dbg_error("Failed to read DH Parms from %s", dfile("dhfile"));
		return;
    }
	parms.size = size;
    res=gnutls_dh_params_import_pkcs3(dh_params, &parms, GNUTLS_X509_FMT_PEM);
    if (res<0) dbg_error("gnutls_dh_params_import_pkcs3");
    dbg_calls("gnutls_dh_params_import_pkcs3");

    parms.data = (uint8_t *)file_load(dfile("rsafile"), &size, MAXDATA);
    if (!parms.data) {
        dbg_error("Failed to read RSA Parms from %s", dfile("rsafile"));
        return;
    }
    parms.size = size;
    res=gnutls_rsa_params_import_pkcs1(rsa_params, &parms, GNUTLS_X509_FMT_PEM);
    if (res<0) dbg_error("gnutls_rsa_params_import_pkcs1");
    dbg_calls("gnutls_rsa_params_import_pkcs1");
}
#endif

void communicate::poll(void) {
	list<channel>::iterator i;
	long ct=time(NULL);
	if (serv_socket_!=-1) {
		struct sockaddr_in cli_addr;
		int clilen=sizeof(cli_addr);
		int cli_socket=accept(serv_socket_, (sockaddr *) &cli_addr, (socklen_t*)&clilen);
		if (cli_socket >= 0) {
			channel ch=channel(1, ct);
			ch.id=channelnr;
			ch.lock=true;
			channelnr++;
			dbg_comm("New %sconnection %i from %s", (type==2?"tls ":""), channelnr, inet_ntoa(cli_addr.sin_addr));
			ch.ip=cli_addr.sin_addr.s_addr;
			ch.handler=serv_handler;
			ch.type=TYPE_SERVER;
			if (type==2) {
                #ifndef no_gnutls
				int res;
				gnutls_session_t session;
				res=gnutls_init (&session, GNUTLS_SERVER);
				if (res<0) dbg_error("gnutls_init %s", gnutls_strerror (res));
				res=gnutls_set_default_export_priority (session);
				if (res<0) dbg_error("gnutls_set_default_export_priority %s", gnutls_strerror (res));
				res=gnutls_credentials_set (session, GNUTLS_CRD_CERTIFICATE, cred);
				if (res<0) dbg_error("gnutls_credentials_set %s", gnutls_strerror (res));
				gnutls_certificate_server_set_request (session, GNUTLS_CERT_REQUEST);
				gnutls_dh_set_prime_bits (session, DH_BITS);

				gnutls_transport_set_ptr (session, (gnutls_transport_ptr_t) cli_socket);
				int ret = gnutls_handshake (session);
				if (ret<0) {
					dbg_error("gnutls_handshake %s", gnutls_strerror (ret));
					close(cli_socket);
					gnutls_deinit(session);
					return;
				}
				ch.stream=new tls_conn(cli_socket, session);
				#endif
			} else {
				ch.stream=new socket_conn(cli_socket);
			}
			set_blocking(cli_socket, false);
			conns_.push_back(ch);
		}
	}
	for(i=conns_.begin(); i!=conns_.end(); i++) {
		i->poll(ct);
		if (i->finished && (serv_socket_!=-1 || type==2)) {
            dbg_comm("Closed connection %i", i->id);
			if (i->type != TYPE_FILE) {
				i->handler.emit(&(*i));
			}
			delete i->stream;
			i=conns_.erase(i);
		}
	}
}

void communicate::complete(int timeout) {
    dbg_calls("Complete comm loop for %d seconds", timeout);
	long current_time=0;
	long start=time(NULL);
	if (serv_socket_!=-1) dbg_fatal("Cannot run a limited loop on a server");
	bool finished=false;
	int count=1;
	while(!finished && count>0) {
		current_time=time(NULL);
		count=0;
		if (current_time > start+timeout) {
		    dbg_comm("Stopped comm loop due to timeout");
		    finished=true;
		}
		list<channel>::iterator i;
		for(i=conns_.begin(); i!=conns_.end(); i++) {
			i->poll(current_time);
			count++;
			if (i->finished || finished) {
			    dbg_comm("Closed connection %i", i->id);
				if (i->type != TYPE_FILE) {
					i->handler.emit(&(*i));
				}
				delete i->stream;
				i=conns_.erase(i);
			}
		}
		usleep(20000);
	}
	dbg_calls("Ready comm");
}

void communicate::server_close(void) {
	list<channel>::iterator i;
	for(i=conns_.begin(); i!=conns_.end(); i++) {
		i->close();
	}
	shutdown(serv_socket_, SHUT_RDWR);
	close(serv_socket_);
}
